import sqlparse
import json
from typing import Dict, List, Any, Optional
from sqlparse.sql import Identifier, Token, Parenthesis
from sqlparse.tokens import Name, Punctuation, Whitespace
import re


class SQLTableParser:
    # 数据库类型到Python类型的映射
    TYPE_MAPPING = {
        "varchar": "str",
        "char": "str",
        "varchar2": "str",
        "nvarchar2": "str",
        "text": "str",
        "clob": "str",
        "nclob": "str",
        "mediumtext": "str",
        "longtext": "str",
        "int": "int",
        "integer": "int",
        "tinyint": "int",
        "smallint": "int",
        "mediumint": "int",
        "bigint": "int",
        "number": "float",
        "numeric": "float",
        "float": "float",
        "double": "float",
        "decimal": "float",
        "real": "float",
        "datetime": "datetime",
        "date": "date",
        "time": "time",
        "timestamp": "datetime",
        "year": "int",
        "boolean": "bool",
        "bool": "bool",
        "blob": "bytes",
        "longblob": "bytes",  
        "raw": "bytes",
        "byte": "bytes",
        "enum": "str",
        "set": "list",
    }

    def parse_sql(self, sql: str) -> Dict[str, Any]:
        """
        解析SQL建表语句

        Args:
            sql: SQL建表语句字符串

        Returns:
            包含表名、字段结构和表注释的字典
        """
        parsed = sqlparse.parse(sql)
        if not parsed:
            raise ValueError("无法解析SQL语句")

        stmt = parsed[0]

        # 提取表名
        table_name = self._extract_table_name(stmt)

        # 提取表注释
        table_comment = self._extract_table_comment(stmt)

        # 提取字段定义
        fields = self._extract_fields(stmt)

        return {"name": table_name, "struts": fields, "desc": table_comment}

    def _extract_table_name(self, stmt) -> str:
        """提取表名"""
        # 查找CREATE TABLE语句后的标识符
        tokens = stmt.tokens
        for i, token in enumerate(tokens):
            if token.match(sqlparse.tokens.DDL, "CREATE"):
                # 查找TABLE关键字
                for j in range(i + 1, len(tokens)):
                    if tokens[j].match(sqlparse.tokens.Keyword, "TABLE"):
                        # 查找表名
                        for k in range(j + 1, len(tokens)):
                            if isinstance(tokens[k], Identifier):
                                return tokens[k].get_real_name()
                            elif tokens[k].ttype is Name:
                                return tokens[k].value
        raise ValueError("无法解析表名")

    def _extract_table_comment(self, stmt) -> str:
        """提取表注释"""
        # 查找COMMENT关键字
        comment = ""
        for token in stmt.tokens:
            if isinstance(token, sqlparse.sql.Comment):
                # 尝试从注释中提取表注释
                comment_text = token.value.lower()
                if "comment" in comment_text and "=" in comment_text:
                    parts = comment_text.split("=")
                    if len(parts) > 1:
                        comment = parts[1].strip().strip("'\"")
                        break

        # 如果没有找到注释，尝试从语句末尾提取
        if not comment:
            sql_lower = str(stmt).lower()
            comment_match = re.search(r"comment\s*=\s*['\"]([^'\"]+)['\"]", sql_lower)
            if comment_match:
                comment = comment_match.group(1)

        return comment

    def _extract_fields(self, stmt) -> List[Dict[str, str]]:
        """提取字段定义"""
        fields = []

        # 查找括号内的内容（字段定义）
        for token in stmt.tokens:
            if isinstance(token, Parenthesis):
                # 解析括号内的字段定义
                field_tokens = token.tokens
                current_field = None

                for field_token in field_tokens:
                    if field_token.ttype is Name or isinstance(field_token, Identifier):
                        # 开始新字段
                        if current_field:
                            fields.append(self._process_field(current_field))

                        current_field = {
                            "name": field_token.value.strip('`"'),
                            "type": "",
                            "comment": "",
                            "raw_tokens": [field_token],
                        }
                    elif current_field:
                        current_field["raw_tokens"].append(field_token)

                # 处理最后一个字段
                if current_field:
                    fields.append(self._process_field(current_field))
                break

        return fields

    def _process_field(self, field_info: Dict) -> Dict[str, str]:
        """处理单个字段信息"""
        field_name = field_info["name"]
        tokens = field_info["raw_tokens"]

        # 提取字段类型
        field_type = ""
        comment = ""

        # 查找类型和注释
        for i, token in enumerate(tokens):
            token_str = str(token).upper()

            # 查找数据类型
            if not field_type and token.ttype is Name and i > 0:
                # 跳过字段名后的第一个标识符（可能是数据类型）
                field_type = token.value.lower()

            # 查找注释
            if "COMMENT" in token_str:
                # 注释通常在COMMENT关键字后面
                for j in range(i + 1, len(tokens)):
                    if tokens[j].ttype is sqlparse.tokens.Literal.String.Single:
                        comment = tokens[j].value.strip("'")
                        break
                break

        # 映射到Python类型
        py_type = self.TYPE_MAPPING.get(field_type, "str")

        return {"field": field_name, "type": py_type, "desc": comment, "rule": ""}

    def to_json(self, sql: str, indent: int = 2) -> str:
        """
        将SQL建表语句转换为JSON字符串

        Args:
            sql: SQL建表语句
            indent: JSON缩进量

        Returns:
            JSON格式的字符串
        """
        result = self.parse_sql(sql)
        return json.dumps(result, ensure_ascii=False, indent=indent)


# 更健壮的版本 - 使用正则表达式辅助解析
class RobustSQLTableParser(SQLTableParser):
    def _extract_fields(self, stmt) -> List[Dict[str, str]]:
        """使用正则表达式辅助提取字段定义"""
        sql_str = str(stmt)

        # 提取括号内的内容
        fields_match = re.search(r"\(([\s\S]*)\)", sql_str)
        if not fields_match:
            return []

        fields_section = fields_match.group(1)

        # 分割字段定义
        fields = []
        field_lines = re.split(r",\s*\n", fields_section)

        for line in field_lines:
            line = line.strip()
            if not line or self._is_constraint_line(line):
                continue

            field = self._parse_field_line(line)
            if field:
                fields.append(field)

        return fields

    def _is_constraint_line(self, line: str) -> bool:
        """判断是否为约束行"""
        line_lower = line.lower()
        constraint_patterns = [
            r"^primary\s+key",
            r"^foreign\s+key",
            r"^unique\s*\(?",
            r"^check\s*\(?",
            r"^constraint\s+",
        ]

        for pattern in constraint_patterns:
            if re.match(pattern, line_lower):
                return True

        return False

    def _parse_field_line(self, line: str) -> Optional[Dict[str, str]]:
        """解析单个字段行"""
        # 提取字段名
        field_name_match = re.match(r'[`"]?(\w+)[`"]?', line)
        if not field_name_match:
            return None

        field_name = field_name_match.group(1)

        # 提取字段类型
        type_match = re.search(r'[`"]?\w+[`"]?\s+(\w+)', line)
        if not type_match:
            return None

        db_type = type_match.group(1).lower()
        py_type = self.TYPE_MAPPING.get(db_type, "str")

        # 提取注释
        comment_match = re.search(r"COMMENT\s+['\"]([^'\"]+)['\"]", line, re.IGNORECASE)
        comment = comment_match.group(1) if comment_match else ""

        return {"field": field_name, "type": py_type, "desc": comment, "rule": ""}





__all__ = ["SQLTableParser", "RobustSQLTableParser"]

# 使用示例
if __name__ == "__main__":
    # MySQL示例
    mysql_sql = """
    CREATE TABLE `tf_vu_hole_task` (
      `TASK_ID` varchar(20) NOT NULL COMMENT '主键',
      `TASK_NAME` varchar(200) DEFAULT NULL COMMENT '任务名称',
      `TASK_FLAG` varchar(2) DEFAULT NULL COMMENT '任务标识 1：周期执行主任务 2：单次执行主任务  3：周期子任务    4：单次复测任务',
      `TASK_STATUS` varchar(2) DEFAULT NULL COMMENT '任务状态 1：待执行 2：执行中 3：已暂停 4：已终止 5：已完成 6：扫描异常',
      `TASK_TYPE` varchar(2) DEFAULT NULL COMMENT '任务类型 1：自动任务 2：人工任务',
      `TASK_SOURCE` varchar(200) DEFAULT NULL,
      `TASK_MAIN_ID` varchar(20) DEFAULT NULL COMMENT '周期性子任务，复测任务关联主任务字段',
      `WORK_ID` varchar(100) DEFAULT NULL COMMENT '关联工单',
      `TASK_TARGET` mediumtext,
      `TASK_DESC` varchar(4000) DEFAULT NULL COMMENT '任务描述',
      `CREATER` varchar(20) DEFAULT NULL COMMENT '任务创建人',
      `CREATE_TIME` datetime DEFAULT NULL,
      `PLAN_TIME` varchar(20) DEFAULT NULL COMMENT '计划执行时长（单位：时）',
      `ACTUAL_START_TIME` datetime DEFAULT NULL,
      `ACTUAL_END_TIME` datetime DEFAULT NULL,
      `SCANNER_TYPE` varchar(2) DEFAULT NULL COMMENT '扫描方式 1：系统自动分配 2：指定扫描器',
      `ADVANCE_FLAG` varchar(2) DEFAULT NULL COMMENT '高级选项配置标识 1：默认配置 2：高级配置',
      `IP_CHECK_FLAG` varchar(2) DEFAULT NULL COMMENT '是否进行IP检测去重  默认否 0：是 1： 否',
      `IP_SCOPE_CHECK` varchar(200) DEFAULT NULL COMMENT '是否进行扫描范围检测 0：是 1：否',
      `EXE_TYPE` varchar(2) DEFAULT NULL COMMENT '执行方式 1：立即执行2：定时执行 3：每日一次 4：每周一次5：每月一次（按日）6：每月一次（按周）',
      `EXE_PARAM` varchar(4000) DEFAULT NULL COMMENT '执行参数',
      `REPORT_FLAG` varchar(2) DEFAULT NULL COMMENT '是否生成报告 默认否 0：是 1：否',
      `report_param` varchar(100) DEFAULT NULL COMMENT '报告参数，多个参数以逗号隔开， 报告类型 0：HTML 1 ：WORD 2：EXECL',
      `HOLE_TYPE` varchar(2) DEFAULT NULL COMMENT '扫描类型 1：系统漏扫 2 ：应用扫描 3：基线检查 4：弱口令检查',
      `SEND_MAIL` varchar(2) DEFAULT NULL COMMENT '是否发送邮件 0：发送 1：不发送',
      `MAILBOX` varchar(100) DEFAULT NULL COMMENT '邮箱',
      `CROSS_VALIDATION` varchar(2) DEFAULT NULL COMMENT '是否交叉验证 0：是 1否',
      `TASK_TIPS` varchar(2000) DEFAULT NULL COMMENT '任务标签，多个标签用逗号隔开',
      `report_status` varchar(2) DEFAULT NULL COMMENT '报告状态 0 待生成  1 生成中 2 已生成',
      `IS_DISPATCH` varchar(2) DEFAULT NULL COMMENT '是否派单：0不派单1派单',
      `WORKORDER_STATUS` varchar(5) DEFAULT NULL COMMENT '工单状态：01草稿02待派单03处理中04挂起05观察中06请求07待验证08终止09归档',
      `TRUST_ID` varchar(32) DEFAULT NULL COMMENT 'TF_SMP_TRUST_INFO表ID',
      `RESPONSE_CODE` varchar(20) DEFAULT '' COMMENT '任务请求返回码',
      `RESPONSE_MSG` varchar(500) DEFAULT '' COMMENT '任务请求返回信息',
      `SCANNER_PARAM` text,
      `REVIEW_TASK_FLAG` varchar(1) DEFAULT '0' COMMENT '复核任务标识 0 否 1是',
      `CHECK_TYPE` varchar(5) DEFAULT NULL COMMENT '交叉验证方式，1-交叉验证，2-组件验证，多个之间用英文逗号连接',
      PRIMARY KEY (`TASK_ID`)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8 COMMENT='扫描任务信息表';
    """

    # Oracle示例
    oracle_sql = """
    CREATE TABLE tf_vu_hole_task (
      TASK_ID varchar2(20) NOT NULL,
      TASK_NAME varchar2(200) DEFAULT NULL,
      TASK_FLAG varchar2(2) DEFAULT NULL,
      TASK_STATUS varchar2(2) DEFAULT NULL,
      TASK_TYPE varchar2(2) DEFAULT NULL,
      TASK_SOURCE varchar2(200) DEFAULT NULL,
      CREATE_TIME DATE DEFAULT NULL
    );
    """

    try:
        parser = RobustSQLTableParser()

        print("解析MySQL SQL:")
        result = parser.to_json(mysql_sql)
        print(result)

        print("\n解析Oracle SQL:")
        result = parser.to_json(oracle_sql)
        print(result)

    except Exception as e:
        print(f"解析错误: {e}")
        import traceback

        traceback.print_exc()
